# vector2chroma.py
# usage: python ./src/vector/vector2chroma.py

import os,sys
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
sys.path.insert(0, project_root)
from config import project_config
import json
import chromadb

# 配置
LOG_FILE = 'tmp_embedding_waiting.txt'
CHROMA_DB_PATH = project_config.chroma_location #  './chroma_db'  # 本地持久化路径，可选
COLLECTION_NAME = project_config.collection_name # 'text_embeddings'

def read_embeddings_log(file_path):
    """逐行读取 embeddings.log 文件，返回字典列表"""
    data = []
    with open(file_path, 'r', encoding='utf-8') as f:
        for line_num, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue  # 跳过空行
            try:
                item = json.loads(line)
                # 验证必要字段
                if 'text' not in item or 'embedding' not in item:
                    print(f"警告：第 {line_num} 行缺少 'text' 或 'embedding' 字段，跳过。")
                    continue
                if not isinstance(item['embedding'], list):
                    print(f"警告：第 {line_num} 行的 'embedding' 不是列表，跳过。")
                    continue
                if len(item['embedding']) != project_config.embedding_dim_size:
                    print(f"警告：第 {line_num} 行的 embedding 维度为 {len(item['embedding'])}，不是 {project_config.embedding_dim_size}，跳过。")
                    continue
                data.append(item)
            except json.JSONDecodeError as e:
                print(f"第 {line_num} 行 JSON 解析错误: {e}")
    return data

def write_to_chroma(data, db_path, collection_name):
    """
    将包含文本和预计算 embedding 的数据写入 Chroma 向量数据库。

    Args:
        data (List[dict]): 数据列表，每个 dict 包含 'text' 和 'embedding' 字段
        db_path (str): Chroma 持久化存储路径
        collection_name (str): 集合名称
    """

    # 创建持久化客户端
    client = chromadb.PersistentClient(path=db_path)

    # 获取或创建集合（不设置 embedding_function，因为我们使用预计算的 embeddings）
    collection = client.get_or_create_collection(
        name=collection_name,
        metadata={
            "dimension": project_config.embedding_dim_size,
            "description": "Collection with precomputed embeddings"
        }
        # 注意：dimension 是自定义元数据，仅用于记录；实际维度由输入 embeddings 决定
    )

    # 提取数据
    ids = []
    documents = []
    embeddings = []

    for idx, item in enumerate(data):
        if not isinstance(item, dict) or 'text' not in item or 'embedding' not in item:
            raise ValueError(f"数据中第 {idx} 条缺少 'text' 或 'embedding' 字段: {item}")

        doc_id = f"doc_{idx}"
        ids.append(doc_id)
        documents.append(str(item['text']))
        embedding = item['embedding']

        # 确保 embedding 是 float 列表
        if not isinstance(embedding, (list, tuple)) or not all(isinstance(x, (int, float)) for x in embedding):
            raise TypeError(f"第 {idx} 条数据的 embedding 必须是数值列表: {embedding}")
        
        embeddings.append(list(embedding))

    # 添加到集合（直接传入预生成的 embeddings）
    collection.add(
        ids=ids,
        embeddings=embeddings,
        documents=documents
    )

    print(f"成功写入 {len(data)} 条记录到 Chroma 集合 '{collection_name}'（路径: {db_path}）")

def main():
    print("开始读取 embeddings.log...")
    data = read_embeddings_log(LOG_FILE)
    print(f"共读取到 {len(data)} 条有效数据。")

    if data:
        print("正在写入 Chroma 数据库...")
        write_to_chroma(data, CHROMA_DB_PATH, COLLECTION_NAME)
        print("完成！")
    else:
        print("没有有效数据可写入。")

if __name__ == '__main__':
    main()